#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
ZIP文件导入模块 - 用于从ZIP压缩包中导入Shapefile数据到PostgreSQL/PostGIS数据库
"""

import os
import sys
import zipfile
import tempfile
import geopandas as gpd
from sqlalchemy import create_engine
from pathlib import Path
import logging
import shutil

# 检查必要的依赖
try:
    import geoalchemy2
except ImportError:
    print("错误: 缺少必要的依赖 'geoalchemy2'")
    print("请运行: pip install geoalchemy2")
    sys.exit(1)

# 导入项目配置
sys.path.append(str(Path(__file__).resolve().parents[3]))
from config import WATER_DATABASE_CONFIG, WATER_MAIN_DATABASE_CONFIG

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)


def import_zip_to_database(zip_path, table_name, db_config=None, if_exists='replace', create_buffer=False, buffer_distance=0.01):
    """
    将ZIP文件中的Shapefile直接导入到数据库，保持原始字段结构
    
    Args:
        zip_path (str): ZIP文件路径
        table_name (str): 目标表名
        db_config (dict, optional): 数据库配置信息，如果为None则使用WATER_DATABASE_CONFIG
        if_exists (str, optional): 如果表已存在，采取的操作。可选值: 'fail', 'replace', 'append'
        create_buffer (bool, optional): 是否创建扩大边界的表
        buffer_distance (float, optional): 边界扩大距离（度）
        
    Returns:
        bool: 导入是否成功
    """
    # 使用默认数据库配置
    db_config = db_config or WATER_DATABASE_CONFIG

    # 创建临时目录
    temp_dir = tempfile.mkdtemp()
    
    try:
        # 创建数据库连接
        connection_string = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
        engine = create_engine(connection_string)
        
        # 提取ZIP文件
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # 查找.shp文件
        shp_files = list(Path(temp_dir).glob('**/*.shp'))
        if not shp_files:
            raise ValueError(f"在ZIP文件 {zip_path} 中未找到Shapefile文件")
        
        # 读取第一个找到的.shp文件
        shp_path = str(shp_files[0])
        gdf = gpd.read_file(shp_path)
        
        # 确保GeoDataFrame使用正确的CRS
        if gdf.crs is None:
            gdf.crs = "EPSG:4326"
        elif gdf.crs != "EPSG:4326":
            gdf = gdf.to_crs("EPSG:4326")
        
        # 将所有字段名转为小写
        gdf.columns = [col.lower() for col in gdf.columns]
        
        # 导入到数据库
        gdf.to_postgis(
            table_name,
            engine,
            if_exists=if_exists,
            index=False,
            schema='public'
        )
        
        logger.info(f"成功导入 {len(gdf)} 条记录到表 {table_name}")
        
        # 如果需要创建扩大边界的表
        if create_buffer:
            buffer_table_name = f"{table_name}_buffer"
            
            # 创建一个新的GeoDataFrame，包含扩大后的几何对象
            buffer_gdf = gdf.copy()
            
            # 对几何对象应用缓冲区
            # 注意：在EPSG:4326坐标系中，0.5度大约是55公里
            buffer_gdf['geometry'] = buffer_gdf.geometry.buffer(buffer_distance)
            
            # 导入扩大边界的表到数据库
            buffer_gdf.to_postgis(
                buffer_table_name,
                engine,
                if_exists=if_exists,
                index=False,
                schema='public'
            )
            
            logger.info(f"成功创建扩大边界表 {buffer_table_name}，边界扩大 {buffer_distance} 度")
        
        return True
    
    except Exception as e:
        logger.error(f"导入数据到数据库失败: {str(e)}")
        raise
    
    finally:
        # 清理临时目录
        shutil.rmtree(temp_dir, ignore_errors=True)


def create_buffer_table_from_existing(table_name, buffer_distance=0.5, db_config=None):
    """
    从已存在的表创建一个边界扩大的表
    
    Args:
        table_name (str): 源表名
        buffer_distance (float): 边界扩大距离（度）
        db_config (dict, optional): 数据库配置信息
        
    Returns:
        bool: 是否成功
    """
    # 使用默认数据库配置
    db_config = db_config or WATER_DATABASE_CONFIG
    
    try:
        # 创建数据库连接
        connection_string = f"postgresql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}"
        engine = create_engine(connection_string)
        
        # 读取源表
        gdf = gpd.read_postgis(f"SELECT * FROM {table_name}", engine, geom_col='geometry')
        
        # 创建缓冲区
        buffer_table_name = f"{table_name}_buffer"
        buffer_gdf = gdf.copy()
        buffer_gdf['geometry'] = buffer_gdf.geometry.buffer(buffer_distance)
        
        # 导入扩大边界的表到数据库
        buffer_gdf.to_postgis(
            buffer_table_name,
            engine,
            if_exists='replace',
            index=False,
            schema='public'
        )
        
        logger.info(f"成功创建扩大边界表 {buffer_table_name}，边界扩大 {buffer_distance} 度")
        return True
        
    except Exception as e:
        logger.error(f"创建扩大边界表失败: {str(e)}")
        raise


if __name__ == "__main__":
    # 示例用法
    # if len(sys.argv) < 3:
    #     print("用法: python zip_import.py <zip文件路径> <表名>")
    #     sys.exit(1)
    #
    # zip_path = sys.argv[1]
    # table_name = sys.argv[2]
    #
    # try:
    #     success = import_zip_to_database(zip_path, table_name)
    #     if success:
    #         print(f"成功导入 {zip_path} 到数据库表 {table_name}")
    #     else:
    #         print(f"导入 {zip_path} 失败")
    # except Exception as e:
    #     print(f"导入过程中发生错误: {str(e)}")
    #     sys.exit(1)
    
    # 导入流域边界数据并创建扩大边界的表
    import_zip_to_database(
        zip_path="/Users/gaoyu/Downloads/水库合并shp.zip",
        table_name="gis_basins",
        db_config=WATER_MAIN_DATABASE_CONFIG,
        create_buffer=True,
        buffer_distance=0.01
    )
    import_zip_to_database(
        zip_path="/Users/gaoyu/Downloads/水库合并shp.zip",
        table_name="gis_basins",
        db_config=WATER_DATABASE_CONFIG,
        create_buffer=True,
        buffer_distance=0.01
    )

    
    # 如果已经导入了数据，可以单独创建扩大边界的表
    # create_buffer_table_from_existing("gis_basins", buffer_distance=0.5)
